from collections import defaultdict
from random import shuffle
import random
from modified_brain import Brain, Area
from parser_table import *

# 我的命名逻辑：非终止符有谓词（pred）和指称词（deno）两种, 谓词根据变元数不同, 分为0, 1, 2, 3元等……
# 一个n元谓词可以填入一个指称词变为n-1元谓词
# 你可以有不同的命名逻辑, 如VP（替代谓词）, NP（替代指称词）等
# Area for stage 2 (components)
DET = "DET"
SUBJ = "SUBJ"
OBJ = "OBJ"
VT = "VT"
VI = "VI"
VC = "VC"
ADJ = "ADJ"
ADV = "ADV"
# 全局最深的pred层数
max_pred_level = 3

# stage2
stage_2_Areas = {
    SUBJ: {},
    DENO: {},
    ADJ: {},
    PRED0: {},
    PRED1: {},
    PRED2: {},
    PRED3: {},
    ADV: {},
    DET: {},
    AN: {},
    DN: {},
    VN: {},
    VA: {},
    S: {},
    VT: {},
    VI: {},
    VC: {},
    CA: {},
    OBJ: {},
    END: {}
}


# 取幂集, 生成正则表达式测试用例用的
def power_set(set1):
    ans = [[]]
    for i in set1:
        l = len(ans)
        for j in range(l):
            t = []
            t.extend(ans[j])
            t.append(i)
            ans.append(t)
    return ans


# 一部分参考代码, 可以在这个基础上完善
class ParserBrain(Brain):
    def __init__(self, p, lexeme_dict=None, all_areas=None, recurrent_areas=None, initial_areas=None,
                 readout_rules=None):
        Brain.__init__(self, p)

        if readout_rules is None:
            readout_rules = {}
        if initial_areas is None:
            initial_areas = []
        if recurrent_areas is None:
            recurrent_areas = []
        if all_areas is None:
            all_areas = []

        self.component_dict = {}
        self.all_areas = all_areas
        self.recurrent_areas = recurrent_areas
        self.initial_areas = initial_areas

        self.fiber_states = defaultdict()
        self.area_states = defaultdict(set)
        self.activated_fibers = defaultdict(set)
        self.readout_rules = readout_rules
        # self.initialize_states()

        self.lexeme_dict = lexeme_dict
        if lexeme_dict is None:
            self.generate_components()

    # 直接生成组件, 用于分步验证; area_index是在同类脑区中的index, area_name是真实name
    def generic_component(self, assembly_index, comp_name, n=1000, k=30, beta=0.1):
        # 比如说连续读两个DET，那么就有area_name = DET_0  DET_1
        area_name = comp_name + f"_{len(stage_2_Areas[comp_name])}"
        # print(area_name)
        stage_2_Areas[comp_name].setdefault(len(stage_2_Areas[comp_name]), area_name)
        # print(stage_2_Areas)
        # print(stage_2_Areas[comp_name])
        # 此时没有刺激, 只能从外部输入
        self.add_explicit_area(area_name, n, k, beta)  # 每个单词是一个脑区?!
        area: Area = self.areas[area_name]
        # 返回的是Assembly对象
        return area.add_random_assembly(assembly_index)  # assembly_index就是单词,返回一个Assembly类

    # 生成词库的方法
    # 可以参考original parser里的lexeme_dict
    def generate_components(self):
        # key: assembly_name
        COMP_DICT = {
            "the": self.generic_component("the", DET),
            "a": self.generic_component("a", DET),
            "dog": self.generic_component("dog", SUBJ),
            "cat": self.generic_component("cat", SUBJ),
            "kid": self.generic_component("kid", SUBJ),
            "ant": self.generic_component("ant", OBJ),
            "mice": self.generic_component("mice", OBJ),
            "banana": self.generic_component("banana", OBJ),
            "chase": self.generic_component("chase", VT),
            "love": self.generic_component("love", VT),
            "bite": self.generic_component("bite", VT), 
            "hate":self.generic_component("hate", VT),
            "find": self.generic_component("find", VT),
            "follow": self.generic_component("follow", VT),
            "cry": self.generic_component("cry", VI),
            "sleep": self.generic_component("sleep", VI),
            "rest": self.generic_component("rest", VI),
            "dance": self.generic_component("dance", VI),
            "smart": self.generic_component("smart", ADJ),
            "big": self.generic_component("big", ADJ),
            "bad": self.generic_component("bad", ADJ),
            "angry":self.generic_component("angry", ADJ),
            "sad":self.generic_component("sad", ADJ),
            "run": self.generic_component("run", VI),   
            "fly": self.generic_component("fly", VI),
            "look": self.generic_component("look", VI),
            "is": self.generic_component("is", VC),
            "was": self.generic_component("was", VC),
            "are": self.generic_component("are", VC),
            "were": self.generic_component("were", VC),
            "quickly": self.generic_component("quickly", ADV),
            "usually": self.generic_component("usually", ADV),
            "suddenly":self.generic_component("suddenly", ADV),
            "quietly":self.generic_component("quietly", ADV),
            "$":self.generic_component("$", END),
        }
        self.component_dict = COMP_DICT

        for key, value in COMP_DICT.items():
            self.all_areas.append(value.area_name)

    # 生成测试用例,根据你的词表,生成所有合法的句子用于测试
    def generate_sentences(self, size=500, output_dir="test_sentences"):
        # 以正则表达式"(Adj)*(Noun)(V)((Adj)*(Noun)+epsilon)"为例
        # 假设*不允许出现重复的字符
        det_list=[]
        adv_list=[]
        adj_list=[]
        subj_list=[]
        obj_list=[]
        vt_list=[]
        vi_list=[]
        vc_list=[]
        for key, value in self.component_dict.items():
            if ADJ in value.area_name:
                adj_list.append(key)
            elif SUBJ in value.area_name:
                subj_list.append(key)
            elif OBJ in value.area_name:
                obj_list.append(key)
            elif VI in value.area_name:
                vi_list.append(key)
            elif VT in value.area_name:
                vt_list.append(key)
            elif DET in value.area_name:
                det_list.append(key)
            elif VC in value.area_name:
                vc_list.append(key)
            elif ADV in value.area_name:
                adv_list.append(key)

        pow_adj = power_set(adj_list)
        pow_adv = power_set(adv_list)
        # 生成名词词组 DET? ADJ* NN(SUBJ/OBJ)
        sp_list = []
        op_list = []
        for det in det_list:
            for noun in subj_list:
                for ap in pow_adj:
                    np = ""
                    # 实现符号","
                    if random.random()<0.5:
                        np += det+' '
                    # 实现符号"*"
                    for i in ap:
                        np += i + ' '
                    np += noun
                    sp_list.append(np)     
        for det in det_list:
            for noun in obj_list:
                for ap in pow_adj:
                    np = ""
                    # 实现符号","
                    if random.random()<0.5:
                        np += det+' '
                    # 实现符号"*"
                    for i in ap:
                        np += i + ' '
                    np += noun
                    op_list.append(np) 
        # 组装句子
        # 主系表
        # (DET)? (ADJ)* SUBJ VC (ADJ)*
        output_list = []
        for vc in vc_list:
            for i in sp_list:
                for adj in pow_adj:
                    adj_=""
                    for a in adj:
                        adj_+=a+' '
                    output_list.append(i+' '+vc+' '+adj_+'$')
        shuffle(output_list)
        output_list = output_list[:150]
        with open(output_dir+"_VC.txt", "w") as f:
            for line in output_list:
                f.write(line + '\n')

        # 主谓宾
        # (DET)? (ADJ)* SUBJ VT (DET)? (ADJ)* OBJ (ADV)*
        output_list = []
        for vt in vt_list:
            for i in sp_list:
                for j in op_list:
                    for adv in pow_adv:
                        adv_=""
                        for a in adv:
                            adv_+=a+' '
                        output_list.append(i+' '+vt+' '+j+' '+adv_+'$')
        shuffle(output_list)
        output_list = output_list[:150]
        with open(output_dir+"_VT.txt", "w") as f:
            for line in output_list:
                f.write(line + '\n')                
        # 主谓
        # DET? ADJ* NN VI ADV*
        output_list = []
        for vi in vi_list:
            for i in sp_list:
                for adv in pow_adv:
                    adv_=""
                    for a in adv:
                        adv_+=a+' '
                    output_list.append(i + ' ' + vi +' '+adv_+'$')
        shuffle(output_list)
        output_list = output_list[:150]
        with open(output_dir+"_VI.txt", "w") as f:
            for line in output_list:
                f.write(line + '\n')  

    # 通过单词获取对应的assembly
    def getAssembly(self, word):
        return self.component_dict[word]

    # 通过area_name来读出对于的area
    def getWord(self,area_name)->Area:
        return self.areas[area_name]

def create_test_library():
    pb = ParserBrain(0.01)
    # 生成测试用例
    pb.generate_sentences()